#include "ground_remove.h"
#include <algorithm>
#include <cmath>
#include <cassert>
#include <iostream>
#include <numeric>

using namespace std;
using pcl::PointXYZI;

/**
 * @brief 过滤距离异常点，仅保留rMin~rMax之间的点
 */
static void filterCloud(const pcl::PointCloud<PointXYZI>::Ptr& cloud, pcl::PointCloud<PointXYZI> &filteredCloud) {
    filteredCloud.clear();
    for (const auto& pt : cloud->points) {
        float distance = std::hypot(pt.x, pt.y);
        if (distance > rMin && distance < rMax) {
            filteredCloud.push_back(pt);
        }
    }
}

/**
 * @brief 计算点的极坐标cell索引
 * @param x,y  点的坐标
 * @param[out] chIdx 角度通道索引
 * @param[out] binIdx 半径bin索引
 */
void ground_remove::getCellIndexFromPoints(float x, float y, int& chIdx, int& binIdx) {
    float angle = atan2(y, x); // 弧度
    if (angle < 0) angle += 2 * M_PI;
    float dist = hypot(x, y);
    chIdx = static_cast<int>(angle / (2 * M_PI / numChannel));
    binIdx = static_cast<int>((dist - rMin) / ((rMax - rMin) / numBin));
    // 边界约束
    if (chIdx < 0) chIdx = 0;
    if (chIdx >= numChannel) chIdx = numChannel - 1;
    if (binIdx < 0) binIdx = 0;
    if (binIdx >= numBin) binIdx = numBin - 1;
}

/**
 * @brief 将点映射到极坐标网格，统计每格的最小z作为候选地面
 */
static void createAndMapPolarGrid(const pcl::PointCloud<PointXYZI>& cloud,
                           std::array<std::array<Cell, numBin>, numChannel>& polarData) {
    // 初始化每个cell
    for(auto& row : polarData) for(auto& c : row) c = Cell();

    for (const auto& pt : cloud.points) {
        int chI, binI;
        ground_remove::getCellIndexFromPoints(pt.x, pt.y, chI, binI);
        if (chI >= 0 && chI < numChannel && binI >= 0 && binI < numBin) {
            polarData[chI][binI].updateMinZ(pt.z);
        }
    }
}

/**
 * @brief 用相邻极格的中值填补非地面但四周都是地面的cell高度
 */
static void applyMedianFilter(std::array<std::array<Cell, numBin>, numChannel>& polarData){
    for(int ch = 1; ch < numChannel-1; ++ch){
        for(int bin = 1; bin < numBin-1; ++bin){
            if(!polarData[ch][bin].isThisGround()) {
                if (   polarData[ch][bin+1].isThisGround()
                    && polarData[ch][bin-1].isThisGround()
                    && polarData[ch+1][bin].isThisGround()
                    && polarData[ch-1][bin].isThisGround() ) {
                    std::vector<float> nbrs = {
                        polarData[ch][bin+1].getHeight(),
                        polarData[ch][bin-1].getHeight(),
                        polarData[ch+1][bin].getHeight(),
                        polarData[ch-1][bin].getHeight()
                    };
                    std::sort(nbrs.begin(), nbrs.end());
                    float median = 0.5f * (nbrs[1] + nbrs[2]);
                    polarData[ch][bin].updateHeight(median);
                    polarData[ch][bin].updateGround();
                }
            }
        }
    }
}

/**
 * @brief 用多邻域均值填补有离群高度值的地面cell
 */
static void outlierFilter(std::array<std::array<Cell, numBin>, numChannel>& polarData){
    for(int ch = 1; ch < numChannel-1; ++ch){
        for(int bin = 1; bin < numBin-2; ++bin){
            if ( polarData[ch][bin].isThisGround()
                && polarData[ch][bin-1].isThisGround()
                && polarData[ch][bin+1].isThisGround()
                && polarData[ch][bin+2].isThisGround() ) {
                float h1 = polarData[ch][bin-1].getHeight();
                float h2 = polarData[ch][bin].getHeight();
                float h3 = polarData[ch][bin+1].getHeight();
                float h4 = polarData[ch][bin+2].getHeight();
                // 如果本cell为孤立极低值，左右邻域有效
                if (h1 != tHmin && h2 == tHmin && h3 != tHmin) {
                    float newH = 0.5f * (h1 + h3);
                    polarData[ch][bin].updateHeight(newH);
                    polarData[ch][bin].updateGround();
                } else if (h1 != tHmin && h2 == tHmin && h3 == tHmin && h4 != tHmin) {
                    float newH = 0.5f * (h1 + h4);
                    polarData[ch][bin].updateHeight(newH);
                    polarData[ch][bin].updateGround();
                }
            }
        }
    }
}

/**
 * @brief 主地面与高点分离方法
 */
void ground_remove::removeGround(
    const pcl::PointCloud<PointXYZI>::Ptr& cloud,
    pcl::PointCloud<PointXYZI>::Ptr& groundCloud,
    pcl::PointCloud<PointXYZI>::Ptr& elevatedCloud) 
{
    // 1. 距离过滤
    pcl::PointCloud<PointXYZI> filteredCloud;
    filterCloud(cloud, filteredCloud);

    // 2. 极坐标网格最小z统计
    std::array<std::array<Cell, numBin>, numChannel> polarData;
    createAndMapPolarGrid(filteredCloud, polarData);

    // 3. 【地面判定主for循环，可通过高度阈值/相邻极格差值等策略】
    // 这里以简单"阈值"示例（可扩展更复杂逻辑）
    for(int ch=0; ch<numChannel; ++ch) {
        for(int bin=0; bin<numBin; ++bin) {
            // 设定全局z阈值或用相邻格差判断
            float minZ = polarData[ch][bin].getHeight();
            if(minZ < 999.0f && minZ > tHmin) {
                polarData[ch][bin].updateGround();
            }
        }
    }

    // 4. 中值+离群补正
    applyMedianFilter(polarData);
    outlierFilter(polarData);

    // 5. 回填点到地面/高点点云
    groundCloud->clear();
    elevatedCloud->clear();
    for(const auto& pt : filteredCloud.points) {
        int ch, bin;
        ground_remove::getCellIndexFromPoints(pt.x, pt.y, ch, bin);
        if(polarData[ch][bin].isThisGround()) 
        {
            float z0 = polarData[ch][bin].getHeight();
            float z1 = (bin>0) ? polarData[ch][bin-1].getHeight() : z0;
            float z2 = (bin<numBin-1) ? polarData[ch][bin+1].getHeight() : z0;
            float z3 = (ch>0) ? polarData[ch-1][bin].getHeight() : z0;
            float z4 = (ch<numChannel-1) ? polarData[ch+1][bin].getHeight() : z0;
            std::vector<float> zs = {z0, z1, z2, z3, z4};
            std::vector<float> valid_zs; 
            for(auto z: zs) if(z < 999 && z > tHmin) valid_zs.push_back(z);
            float refZ = valid_zs.empty() ? z0 : std::accumulate(valid_zs.begin(), valid_zs.end(),0.f)/valid_zs.size();
            if(std::abs(pt.z - refZ) < 0.3f) { // <--- 阈值可调
                groundCloud->push_back(pt);
            } else {
                elevatedCloud->push_back(pt);
            }
        } else {
            elevatedCloud->push_back(pt);
        }
    }
}